{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "d7caa42e",
   "metadata": {},
   "source": [
    "# Tutorial 03: weak imposition of Dirichlet BCs by a Lagrange multiplier (linear problem)\n",
    "\n",
    "In this tutorial we solve the problem\n",
    "\n",
    "$$\\begin{cases}\n",
    "-\\Delta u = f, & \\text{in } \\Omega,\\\\\n",
    " u   = g, & \\text{on } \\Gamma = \\partial\\Omega,\n",
    "\\end{cases}$$\n",
    "\n",
    "where $\\Omega$ is a ball in 2D.\n",
    "\n",
    "We compare the following two cases:\n",
    "* **strong imposition of Dirichlet BCs**:\n",
    "the corresponding weak formulation is\n",
    "$$\n",
    "\\text{find } u \\in V_g \\text{ s.t. } \\int_\\Omega \\nabla u \\cdot \\nabla v \\ dx = \\int_\\Omega f v \\ dx, \\quad \\forall v \\in V_0\\\\\n",
    "$$\n",
    "where\n",
    "$$\n",
    "V_g = \\{v \\in H^1(\\Omega): v|_\\Gamma = g\\},\\\\\n",
    "V_0 = \\{v \\in H^1(\\Omega): v|_\\Gamma = 0\\}.\\\\\n",
    "$$\n",
    "* **weak imposition of Dirichlet BCs**: this requires an introduction of a multiplier $\\lambda$ which is restricted to $\\Gamma$, and solves\n",
    "$$\n",
    "\\text{find } w, \\lambda \\in V \\times M \\text{ s.t. }\\\\\n",
    "\n",
    "\\begin{align*}\n",
    "&\\int_\\Omega \\nabla w \\cdot \\nabla v \\ dx & &+ \\int_\\Gamma \\lambda v \\ ds & &= \\int_\\Omega f v \\ dx, & \\forall v \\in V,\\\\\n",
    "&\\int_\\Gamma w \\mu \\ ds & & & &= \\int_\\Gamma g \\mu \\ ds, & \\forall \\mu \\in M\n",
    "\n",
    "\\end{align*}\n",
    "$$\n",
    "where\n",
    "$$\n",
    "V = H^1(\\Omega),\\\\\n",
    "M = L^{2}(\\Gamma).\\\\\n",
    "$$\n",
    "\n",
    "This example is a prototypical case of problems containing subdomain/boundary restricted variables (the Lagrange multiplier, in this case)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "58b38cc2",
   "metadata": {},
   "outputs": [],
   "source": [
    "import dolfinx.fem\n",
    "import dolfinx.fem.petsc\n",
    "import dolfinx.io\n",
    "import gmsh\n",
    "import mpi4py.MPI\n",
    "import numpy as np\n",
    "import petsc4py.PETSc\n",
    "import ufl\n",
    "import viskex"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9b9b2a4d",
   "metadata": {},
   "outputs": [],
   "source": [
    "import multiphenicsx.fem\n",
    "import multiphenicsx.fem.petsc"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6941da68",
   "metadata": {},
   "source": [
    "### Geometrical parameters"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e6fe6d4b",
   "metadata": {},
   "outputs": [],
   "source": [
    "r = 3\n",
    "mesh_size = 1. / 4."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c75137ab",
   "metadata": {},
   "source": [
    "### Mesh"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "df34d60c",
   "metadata": {},
   "outputs": [],
   "source": [
    "gmsh.initialize()\n",
    "gmsh.model.add(\"mesh\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1d6029dd",
   "metadata": {},
   "outputs": [],
   "source": [
    "p0 = gmsh.model.geo.addPoint(0.0, 0.0, 0.0, mesh_size)\n",
    "p1 = gmsh.model.geo.addPoint(0.0, +r, 0.0, mesh_size)\n",
    "p2 = gmsh.model.geo.addPoint(0.0, -r, 0.0, mesh_size)\n",
    "c0 = gmsh.model.geo.addCircleArc(p1, p0, p2)\n",
    "c1 = gmsh.model.geo.addCircleArc(p2, p0, p1)\n",
    "boundary = gmsh.model.geo.addCurveLoop([c0, c1])\n",
    "domain = gmsh.model.geo.addPlaneSurface([boundary])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c453e960",
   "metadata": {},
   "outputs": [],
   "source": [
    "gmsh.model.geo.synchronize()\n",
    "gmsh.model.addPhysicalGroup(1, [c0, c1], 1)\n",
    "gmsh.model.addPhysicalGroup(2, [boundary], 0)\n",
    "gmsh.model.mesh.generate(2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f5209352",
   "metadata": {},
   "outputs": [],
   "source": [
    "mesh, subdomains, boundaries = dolfinx.io.gmshio.model_to_mesh(\n",
    "    gmsh.model, comm=mpi4py.MPI.COMM_WORLD, rank=0, gdim=2)\n",
    "gmsh.finalize()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8f680d79-bc3d-47c2-a739-dc833bd37d08",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create connectivities required by the rest of the code\n",
    "mesh.topology.create_connectivity(mesh.topology.dim - 1, mesh.topology.dim)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "aac7a51a",
   "metadata": {},
   "outputs": [],
   "source": [
    "facets_Gamma = boundaries.indices[boundaries.values == 1]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5f9ddf72",
   "metadata": {},
   "outputs": [],
   "source": [
    "viskex.dolfinx.plot_mesh(mesh)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "46bc298a",
   "metadata": {},
   "outputs": [],
   "source": [
    "viskex.dolfinx.plot_mesh_tags(mesh, boundaries, \"boundaries\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "be4a9020",
   "metadata": {},
   "source": [
    "### Weak imposition of Dirichlet BCs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e6f13596",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Define a function space\n",
    "V = dolfinx.fem.functionspace(mesh, (\"Lagrange\", 2))\n",
    "M = V.clone()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "deef1cac",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Define restrictions\n",
    "dofs_V = np.arange(0, V.dofmap.index_map.size_local + V.dofmap.index_map.num_ghosts)\n",
    "dofs_M_Gamma = dolfinx.fem.locate_dofs_topological(M, boundaries.dim, facets_Gamma)\n",
    "restriction_V = multiphenicsx.fem.DofMapRestriction(V.dofmap, dofs_V)\n",
    "restriction_M_Gamma = multiphenicsx.fem.DofMapRestriction(M.dofmap, dofs_M_Gamma)\n",
    "restriction = [restriction_V, restriction_M_Gamma]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fc839d3d",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Define trial and test functions\n",
    "(u, l) = (ufl.TrialFunction(V), ufl.TrialFunction(M))\n",
    "(v, m) = (ufl.TestFunction(V), ufl.TestFunction(M))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9e80eb0a",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Define problem block forms\n",
    "g = dolfinx.fem.Function(V)\n",
    "g.interpolate(lambda x: np.sin(3 * x[0] + 1) * np.sin(3 * x[1] + 1))\n",
    "a = [[ufl.inner(ufl.grad(u), ufl.grad(v)) * ufl.dx, ufl.inner(l, v) * ufl.ds],\n",
    "     [ufl.inner(u, m) * ufl.ds, None]]\n",
    "f = [ufl.inner(1, v) * ufl.dx, ufl.inner(g, m) * ufl.ds]\n",
    "a_cpp = dolfinx.fem.form(a)\n",
    "f_cpp = dolfinx.fem.form(f)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4b7e670e",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Assemble the block linear system\n",
    "A = multiphenicsx.fem.petsc.assemble_matrix_block(a_cpp, bcs=[], restriction=(restriction, restriction))\n",
    "A.assemble()\n",
    "F = multiphenicsx.fem.petsc.assemble_vector_block(f_cpp, a_cpp, bcs=[], restriction=restriction)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e4df149b",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Solve\n",
    "ul = multiphenicsx.fem.petsc.create_vector_block(f_cpp, restriction=restriction)\n",
    "ksp = petsc4py.PETSc.KSP()\n",
    "ksp.create(mesh.comm)\n",
    "ksp.setOperators(A)\n",
    "ksp.setType(\"preonly\")\n",
    "ksp.getPC().setType(\"lu\")\n",
    "ksp.getPC().setFactorSolverType(\"mumps\")\n",
    "ksp.setFromOptions()\n",
    "ksp.solve(F, ul)\n",
    "ul.ghostUpdate(addv=petsc4py.PETSc.InsertMode.INSERT, mode=petsc4py.PETSc.ScatterMode.FORWARD)\n",
    "ksp.destroy()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "21d6de16",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Split the block solution in components\n",
    "(u, l) = (dolfinx.fem.Function(V), dolfinx.fem.Function(M))\n",
    "with multiphenicsx.fem.petsc.BlockVecSubVectorWrapper(ul, [V.dofmap, M.dofmap], restriction) as ul_wrapper:\n",
    "    for ul_wrapper_local, component in zip(ul_wrapper, (u, l)):\n",
    "        with component.x.petsc_vec.localForm() as component_local:\n",
    "            component_local[:] = ul_wrapper_local\n",
    "ul.destroy()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "58f61dc9",
   "metadata": {},
   "outputs": [],
   "source": [
    "viskex.dolfinx.plot_scalar_field(u, \"u\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "97e090c0",
   "metadata": {},
   "outputs": [],
   "source": [
    "viskex.dolfinx.plot_scalar_field(l, \"l\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "27d28dca",
   "metadata": {},
   "source": [
    "### Strong imposition of Dirichlet BCs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "78fcb18e",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Define Dirichlet BC object on Gamma\n",
    "dofs_V_Gamma = dolfinx.fem.locate_dofs_topological(V, boundaries.dim, facets_Gamma)\n",
    "bc_ex = dolfinx.fem.dirichletbc(g, dofs_V_Gamma)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "306b95a1",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Solve\n",
    "u_ex = dolfinx.fem.Function(V)\n",
    "problem_ex = dolfinx.fem.petsc.LinearProblem(\n",
    "    a[0][0], f[0], bcs=[bc_ex], u=u_ex,\n",
    "    petsc_options={\"ksp_type\": \"preonly\", \"pc_type\": \"lu\", \"pc_factor_mat_solver_type\": \"mumps\"})\n",
    "problem_ex.solve()\n",
    "u_ex.x.petsc_vec.ghostUpdate(addv=petsc4py.PETSc.InsertMode.INSERT, mode=petsc4py.PETSc.ScatterMode.FORWARD)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "937ad93f",
   "metadata": {},
   "outputs": [],
   "source": [
    "viskex.dolfinx.plot_scalar_field(u_ex, \"u\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d2795bdb",
   "metadata": {},
   "source": [
    "### Comparison and error computation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "32383aed",
   "metadata": {},
   "outputs": [],
   "source": [
    "u_ex_norm = np.sqrt(mesh.comm.allreduce(\n",
    "    dolfinx.fem.assemble_scalar(dolfinx.fem.form(ufl.inner(ufl.grad(u_ex), ufl.grad(u_ex)) * ufl.dx)),\n",
    "    op=mpi4py.MPI.SUM))\n",
    "err_norm = np.sqrt(mesh.comm.allreduce(\n",
    "    dolfinx.fem.assemble_scalar(dolfinx.fem.form(ufl.inner(ufl.grad(u_ex - u), ufl.grad(u_ex - u)) * ufl.dx)),\n",
    "    op=mpi4py.MPI.SUM))\n",
    "print(\"Relative error is equal to\", err_norm / u_ex_norm)\n",
    "assert np.isclose(err_norm / u_ex_norm, 0., atol=1.e-10)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython"
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
